from app import create_app
from models import db
from sqlalchemy import text

def migrate_database():
    app = create_app()
    
    with app.app_context():
        # 检查并添加缺失的字段
        try:
            # 检查users表是否有created_at字段
            result = db.session.execute(text("SHOW COLUMNS FROM users LIKE 'created_at'"))
            if not result.fetchone():
                print("添加users表的created_at字段...")
                db.session.execute(text("ALTER TABLE users ADD COLUMN created_at DATETIME DEFAULT CURRENT_TIMESTAMP"))
                db.session.commit()
                print("users表的created_at字段已添加")
            else:
                print("users表的created_at字段已存在")
                
            # 检查courses表是否有created_at字段
            result = db.session.execute(text("SHOW COLUMNS FROM courses LIKE 'created_at'"))
            if not result.fetchone():
                print("添加courses表的created_at字段...")
                db.session.execute(text("ALTER TABLE courses ADD COLUMN created_at DATETIME DEFAULT CURRENT_TIMESTAMP"))
                db.session.commit()
                print("courses表的created_at字段已添加")
            else:
                print("courses表的created_at字段已存在")
                
            # 检查problems表是否有created_at字段
            result = db.session.execute(text("SHOW COLUMNS FROM problems LIKE 'created_at'"))
            if not result.fetchone():
                print("添加problems表的created_at字段...")
                db.session.execute(text("ALTER TABLE problems ADD COLUMN created_at DATETIME DEFAULT CURRENT_TIMESTAMP"))
                db.session.commit()
                print("problems表的created_at字段已添加")
            else:
                print("problems表的created_at字段已存在")
                
            # 检查homeworks表是否有created_at字段
            result = db.session.execute(text("SHOW COLUMNS FROM homeworks LIKE 'created_at'"))
            if not result.fetchone():
                print("添加homeworks表的created_at字段...")
                db.session.execute(text("ALTER TABLE homeworks ADD COLUMN created_at DATETIME DEFAULT CURRENT_TIMESTAMP"))
                db.session.commit()
                print("homeworks表的created_at字段已添加")
            else:
                print("homeworks表的created_at字段已存在")
                
            print("数据库迁移完成！")
            
        except Exception as e:
            print(f"迁移过程中出现错误: {e}")
            db.session.rollback()

if __name__ == '__main__':
    migrate_database()
